# 自定义 MemorySaver 类，限制最多保存 5 轮对话

from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata
from langgraph.checkpoint.memory import MemorySaver


# 定义检查点数据结构

# 修改 LimitedMemorySaver 类中的 put 方法如下：
class LimitedMemorySaver(MemorySaver):
    def __init__(self, max_turns: int = 5):
        super().__init__()
        self.max_turns = max_turns

    def put(self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, **kwargs) -> None:
        # 获取当前会话的 thread_id
        thread_id = config["configurable"]["thread_id"]
        # 获取现有检查点
        existing = self.get(config) or {"messages": []}
        messages = existing["messages"]
        # 添加新消息
        messages.extend(checkpoint["messages"])
        # 计算轮数（每轮包含 HumanMessage 和 AIMessage）
        turns = [msg for msg in messages if isinstance(msg, (HumanMessage, AIMessage))]
        # 保留最后 N 轮（每轮2条消息）
        if len(turns) > self.max_turns * 2:
            messages = messages[-(self.max_turns * 2):]
            checkpoint["messages"] = messages
        # 调用父类方法保存
        super().put(config, checkpoint, metadata, **kwargs)


